package org.jtzc.springaiembedding.controller;

import org.jtzc.springaiembedding.util.Util;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * @author wu chuang
 * @description
 */
@RestController
public class EmbeddingController {

    private final EmbeddingModel embeddingModel;

    @Autowired
    public EmbeddingController(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    @Autowired
    private DocumentTransformer tokenTextSplitter;
    @Autowired
    private VectorStore vectorStore;

    @Value("classpath:CV.pdf")
    private Resource resource;

    @GetMapping("/init")
    public void addDocumentToVectorDB() {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
        List<Document> fileDocuments = tikaDocumentReader.get();
        List<Document> documents = tokenTextSplitter.apply(fileDocuments);
        vectorStore.accept(documents);

    }


    @GetMapping("/vectorQuery2")
    public List<String> vectorQuery2(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        List<Document> documents = vectorStore.similaritySearch(message);
        // 获取每个Document里的content
        List<String> collect = documents.stream().map(Document::getFormattedContent).toList();
        return collect;
    }

    @GetMapping("/vectorQuery")
    public List<Document> vectorQuery() {
        /**
         * Map.of(
         *     "source", "https://example.com",  // 文档来源
         *     "author", "Jane Doe",             // 作者信息
         *     "timestamp", Instant.now()        // 创建时间戳
         * )
         */
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! ...", Map.of("meta1", "meta1")),
                new Document("The World is Big...", Collections.emptyMap()),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));

        vectorStore.add(documents);

        List<Document> results = vectorStore.similaritySearch(SearchRequest.builder()
                .query("Spring")
                .topK(5)
                .build());

        return results;
    }

    @GetMapping("/vectorQuery3")
    public String vectorQuery3() {
        /**
         * Map.of(
         *     "source", "https://example.com",  // 文档来源
         *     "author", "Jane Doe",             // 作者信息
         *     "timestamp", Instant.now()        // 创建时间戳
         * )
         */
        List<Document> documents = List.of(
                new Document("OpenAI的ChatGPT是一个强大的语言模型。", Map.of("meta1", "meta1")),
                new Document("天空是蓝色的,阳光灿烂。", Collections.emptyMap()),
                new Document("人工智能正在改变世界。", Map.of("meta2", "meta2")),
                new Document("java是一种流行的编程语言。", Map.of("meta2", "meta2")));

        vectorStore.add(documents);
        String message = "天空是什么颜色的？";
        EmbeddingResponse query_embedding = this.embeddingModel.embedForResponse(List.of(message));
        documents.forEach(item -> {
            EmbeddingResponse document_embeddings = this.embeddingModel.embedForResponse(List.of(item.getFormattedContent()));
            float v = Util.cosineSimilarity(document_embeddings.getResult().getOutput(), query_embedding.getResult().getOutput());
            System.out.println(v);
            System.out.println(item.getFormattedContent());
        });
        return "results";
    }



    @GetMapping("/embedding")
    public Map<String,Object> embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(message,"hell world"));
        System.out.println(embeddingResponse.getResults().size());
        System.out.println(embeddingResponse.getResults().get(0).getOutput().length);
        return Map.of("embedding", embeddingResponse);
    }
    @GetMapping("/embedding2")
    public Map<String,Object> embed2(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        EmbeddingRequest embeddingRequest =
                new EmbeddingRequest(List.of(message),
                        OpenAiEmbeddingOptions.builder()
                                .model("text-embedding-ada-002")
                                .dimensions(3) //维度
                                .build());

        EmbeddingResponse embeddingResponse = embeddingModel.call(embeddingRequest);
        return Map.of("embedding", embeddingResponse);
    }
    @GetMapping("/vectorQuery22")
    public List<String> vectorQuery2() {
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
                .query("竞赛经历")
                .topK(2)
                .build());
        List<String> collect = documents.stream().map(Document::getFormattedContent).toList();
        collect.forEach(e->{
            System.out.println("-------");
            System.out.println(e);
        });
        return collect;
    }


}
